import adaboost
import numpy as np
import time

DATASET_PATH = '../DataSet/mnist/'

train_images = np.load(DATASET_PATH + 'train_img.npy')
train_labels = np.load(DATASET_PATH + 'train_label.npy')

test_images = np.load(DATASET_PATH + 'test_img.npy')
test_labels = np.load(DATASET_PATH + 'test_label.npy')

print(train_labels.shape)

my_trainset = [train_images[i].reshape(-1) for i in range(60000) if train_labels[i][1] == 1 or train_labels[i][9] == 1]
my_labels = [train_labels[i][9] for i in range(60000) if train_labels[i][1] == 1 or train_labels[i][9] == 1]

my_trainset = np.array(my_trainset, np.float) / 255 - 0.5
my_labels = np.array(my_labels)
my_labels[my_labels <= 0] = -1
my_labels[my_labels > 0] = 1

print(my_trainset.shape)
print(my_labels.shape)

train_x, train_y = my_trainset[:1000], my_labels[:1000]
test_x, test_y = my_trainset[1000:2000], my_labels[1000:2000]

# import matplotlib.pyplot as plt
#
# id = 29
# plt.imshow(train_x[id].reshape((28, 28)))
# plt.show()
# print(train_y[id])

# 训练手写数字1和9的二分类分类器
ada = adaboost.AdaBoost()
ada.train(train_x, train_y)

print('weak-classifier number:', len(ada.classfierArr))
print()

t0 = time.time()
error_count = 0

pred = ada.predict(test_x)
error_count = np.sum(test_y != pred)

print('the test error rate is: %.2f' % (1.0 * error_count / test_y.shape[0]))
print('predict time:%.2fs for %d samples' % (time.time() - t0, test_y.shape[0]))

